import tkinter as tk
from PIL import ImageGrab, Image
import os

class ScreenshotTool:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("截屏工具")
        self.root.geometry("400x200")
        
        # 当前目录作为存储路径
        self.save_dir = os.getcwd()  # 获取当前工作目录
        
        # GUI组件
        self.label = tk.Label(self.root, text="截屏工具 v1.0")
        self.label.pack(pady=10)
        
        self.screenshot_btn = tk.Button(self.root, text="选择区域截图", command=self.start_selection)
        self.screenshot_btn.pack(pady=10)
        
        self.status = tk.Label(self.root, text="")
        self.status.pack(pady=10)
        
        # 选择区域的变量
        self.start_x = None
        self.start_y = None
        self.rect = None
        self.canvas = None
        
        self.root.mainloop()

    def get_unique_filename(self):
        """生成唯一的文件名，避免覆盖"""
        base_name = "screenshot"
        extension = ".png"
        counter = 0
        filename = os.path.join(self.save_dir, f"{base_name}{extension}")
        
        while os.path.exists(filename):
            counter += 1
            filename = os.path.join(self.save_dir, f"{base_name}_{counter}{extension}")
        
        return filename

    def start_selection(self):
        """开始选择截图区域"""
        self.root.withdraw()  # 隐藏主窗口
        
        # 创建全屏透明窗口用于选择区域
        self.selection_window = tk.Toplevel()
        self.selection_window.attributes('-fullscreen', True)
        self.selection_window.attributes('-alpha', 0.3)  # 设置透明度
        
        self.canvas = tk.Canvas(self.selection_window, cursor="cross")
        self.canvas.pack(fill="both", expand=True)
        
        # 绑定鼠标事件
        self.canvas.bind("<ButtonPress-1>", self.on_press)
        self.canvas.bind("<B1-Motion>", self.on_drag)
        self.canvas.bind("<ButtonRelease-1>", self.on_release)
        
        # 按ESC键取消
        self.selection_window.bind("<Escape>", lambda e: self.cancel_selection())

    def on_press(self, event):
        """鼠标按下，开始记录起点"""
        self.start_x = self.canvas.winfo_pointerx()
        self.start_y = self.canvas.winfo_pointery()
        self.rect = self.canvas.create_rectangle(self.start_x, self.start_y, self.start_x, self.start_y, outline="red", width=2)

    def on_drag(self, event):
        """鼠标拖动，更新矩形区域"""
        cur_x = self.canvas.winfo_pointerx()
        cur_y = self.canvas.winfo_pointery()
        self.canvas.coords(self.rect, self.start_x, self.start_y, cur_x, cur_y)

    def on_release(self, event):
        """鼠标释放，完成截图"""
        end_x = self.canvas.winfo_pointerx()
        end_y = self.canvas.winfo_pointery()
        
        # 计算截图区域
        left = min(self.start_x, end_x)
        top = min(self.start_y, end_y)
        right = max(self.start_x, end_x)
        bottom = max(self.start_y, end_y)
        
        # 关闭选择窗口并截图
        self.selection_window.destroy()
        self.take_screenshot(left, top, right, bottom)

    def take_screenshot(self, left, top, right, bottom):
        """根据选择区域截图"""
        try:
            screenshot = ImageGrab.grab(bbox=(left, top, right, bottom))
            save_path = self.get_unique_filename()  # 获取唯一文件名
            screenshot.save(save_path)
            screenshot.close()
            
            self.root.deiconify()  # 恢复主窗口
            self.status.config(text=f"截图已保存为 {os.path.basename(save_path)}")
        except Exception as e:
            self.root.deiconify()
            self.status.config(text=f"截图失败: {str(e)}")

    def cancel_selection(self):
        """取消选择"""
        self.selection_window.destroy()
        self.root.deiconify()
        self.status.config(text="已取消截图")

if __name__ == "__main__":
    ScreenshotTool()